#include <mcp/server/server_session.h>
#include <mcp/server/tool_manager.h>
#include <mcp/server/resource_manager.h>
#include <mcp/server/prompt_manager.h>
#include <mcp/transport/transport.h>
#include <mcp/json/value.h>
#include <mcp/error.h>
#include <mcp/types.h>
#include <iostream>
#include <sstream>
#include <memory>
#include <algorithm>
#include <random>
#include <chrono>

namespace mcp {
namespace server {

// ServerSession implementation
struct ServerSession::Impl {
    std::unique_ptr<transport::Transport> transport;
    std::shared_ptr<ToolManager> tool_manager;
    std::shared_ptr<ResourceManager> resource_manager;
    std::shared_ptr<PromptManager> prompt_manager;

    // Session state
    InitializationState init_state = InitializationState::NotInitialized;
    std::optional<types::InitializeRequestParams> client_params;
    InitializationOptions init_options;

    // Message handling
    std::queue<ServerRequest> incoming_requests;
    std::queue<ServerNotification> incoming_notifications;
    mutable std::mutex message_mutex;
    std::condition_variable message_cv;

    // Request tracking
    std::unordered_map<std::string, std::promise<types::ServerResult>> pending_requests;
    mutable std::mutex request_mutex;
    std::atomic<int64_t> next_request_id{1};

    // Close handler
    std::function<void()> close_handler;

    // Logging
    std::string session_id;
    std::chrono::steady_clock::time_point start_time;

    Impl(std::unique_ptr<transport::Transport> transport,
         std::shared_ptr<ToolManager> tool_manager,
         std::shared_ptr<ResourceManager> resource_manager,
         std::shared_ptr<PromptManager> prompt_manager,
         const InitializationOptions& init_options)
        : transport(std::move(transport))
        , tool_manager(std::move(tool_manager))
        , resource_manager(std::move(resource_manager))
        , prompt_manager(std::move(prompt_manager))
        , init_options(init_options)
        , start_time(std::chrono::steady_clock::now()) {

        // Generate session ID
        std::random_device rd;
        std::mt19937 gen(rd());
        std::uniform_int_distribution<uint64_t> dis;
        std::stringstream ss;
        ss << std::hex << dis(gen);
        session_id = ss.str();
    }

    std::string generate_request_id() {
        return std::to_string(next_request_id++);
    }
};

ServerSession::ServerSession(std::unique_ptr<transport::Transport> transport,
                           std::shared_ptr<ToolManager> tool_manager,
                           std::shared_ptr<ResourceManager> resource_manager,
                           std::shared_ptr<PromptManager> prompt_manager,
                           const InitializationOptions& init_options)
    : impl_(std::make_unique<Impl>(std::move(transport), std::move(tool_manager),
                                  std::move(resource_manager), std::move(prompt_manager), init_options)) {
}

ServerSession::~ServerSession() = default;

// Session management
void ServerSession::start() {
    if (!impl_->transport) {
        throw std::runtime_error("No transport available");
    }

    // Start the transport message processing loop
    impl_->transport->start([this](const json::Value& message) {
        handle_message(message);
    });
}

void ServerSession::stop() {
    if (impl_->transport) {
        impl_->transport->stop();
    }

    // Notify all pending requests that session is closing
    std::lock_guard<std::mutex> lock(impl_->request_mutex);
    for (auto& [id, promise] : impl_->pending_requests) {
        try {
            promise.set_value(types::ServerResult{Error::internal_error("Session closed")});
        } catch (...) {
            // Promise might already be satisfied
        }
    }
    impl_->pending_requests.clear();
}

bool ServerSession::is_connected() const {
    return impl_->transport && impl_->transport->is_connected();
}

InitializationState ServerSession::initialization_state() const {
    return impl_->init_state;
}

std::optional<types::InitializeRequestParams> ServerSession::client_params() const {
    return impl_->client_params;
}

// Capabilities checking
bool ServerSession::check_client_capability(const types::ClientCapabilities& capability) const {
    if (!impl_->client_params) {
        return false;
    }

    const auto& client_caps = impl_->client_params->capabilities;

    // Check roots capability
    if (capability.roots) {
        if (!client_caps.roots) {
            return false;
        }
        if (capability.roots->list_changed && !client_caps.roots->list_changed) {
            return false;
        }
    }

    // Check sampling capability
    if (capability.sampling && !client_caps.sampling) {
        return false;
    }

    // Check elicitation capability
    if (capability.elicitation && !client_caps.elicitation) {
        return false;
    }

    // Check experimental capabilities
    if (capability.experimental) {
        if (!client_caps.experimental) {
            return false;
        }
        for (const auto& [key, value] : *capability.experimental) {
            auto it = client_caps.experimental->find(key);
            if (it == client_caps.experimental->end() || it->second != value) {
                return false;
            }
        }
    }

    return true;
}

// Message handling
void ServerSession::handle_message(const json::Value& message) {
    try {
        if (!message.is_object()) {
            throw std::runtime_error("Message must be an object");
        }

        std::string method = message.get("method", "").as_string();
        bool has_id = message.contains("id");

        if (has_id) {
            // This is a request
            ServerRequest request;
            request.id = message["id"].as_string();
            request.method = method;
            request.params = message.get("params", json::Value::object({}));

            std::lock_guard<std::mutex> lock(impl_->message_mutex);
            impl_->incoming_requests.push(std::move(request));
            impl_->message_cv.notify_one();
        } else {
            // This is a notification
            ServerNotification notification;
            notification.method = method;
            notification.params = message.get("params", json::Value::object({}));

            std::lock_guard<std::mutex> lock(impl_->message_mutex);
            impl_->incoming_notifications.push(std::move(notification));
            impl_->message_cv.notify_one();
        }
    } catch (const std::exception& e) {
        // Log error but continue processing
        std::cerr << "Error handling message: " << e.what() << std::endl;
    }
}

std::optional<ServerRequest> ServerSession::receive_request(std::chrono::milliseconds timeout) {
    std::unique_lock<std::mutex> lock(impl_->message_mutex);

    if (impl_->incoming_requests.empty()) {
        if (timeout.count() > 0) {
            impl_->message_cv.wait_for(lock, timeout, [this] {
                return !impl_->incoming_requests.empty();
            });
        } else {
            impl_->message_cv.wait(lock, [this] {
                return !impl_->incoming_requests.empty();
            });
        }
    }

    if (impl_->incoming_requests.empty()) {
        return std::nullopt;
    }

    auto request = std::move(impl_->incoming_requests.front());
    impl_->incoming_requests.pop();
    return request;
}

std::optional<ServerNotification> ServerSession::receive_notification(std::chrono::milliseconds timeout) {
    std::unique_lock<std::mutex> lock(impl_->message_mutex);

    if (impl_->incoming_notifications.empty()) {
        if (timeout.count() > 0) {
            impl_->message_cv.wait_for(lock, timeout, [this] {
                return !impl_->incoming_notifications.empty();
            });
        } else {
            impl_->message_cv.wait(lock, [this] {
                return !impl_->incoming_notifications.empty();
            });
        }
    }

    if (impl_->incoming_notifications.empty()) {
        return std::nullopt;
    }

    auto notification = std::move(impl_->incoming_notifications.front());
    impl_->incoming_notifications.pop();
    return notification;
}

// Request/response handling
std::future<types::ServerResult> ServerSession::send_request(const ServerRequest& request) {
    auto promise = std::make_shared<std::promise<types::ServerResult>>();
    auto future = promise->get_future();

    std::string request_id = impl_->generate_request_id();

    // Store the promise for later fulfillment
    {
        std::lock_guard<std::mutex> lock(impl_->request_mutex);
        impl_->pending_requests[request_id] = std::move(*promise);
    }

    // Send the request
    json::Value message = json::Value::object({
        {"jsonrpc", "2.0"},
        {"id", request_id},
        {"method", request.method},
        {"params", request.params}
    });

    if (impl_->transport) {
        impl_->transport->send_message(message);
    } else {
        // If no transport, fulfill with error
        std::lock_guard<std::mutex> lock(impl_->request_mutex);
        auto it = impl_->pending_requests.find(request_id);
        if (it != impl_->pending_requests.end()) {
            it->second.set_value(types::ServerResult{Error::internal_error("No transport available")});
            impl_->pending_requests.erase(it);
        }
    }

    return future;
}

void ServerSession::send_notification(const ServerNotification& notification) {
    json::Value message = json::Value::object({
        {"jsonrpc", "2.0"},
        {"method", notification.method},
        {"params", notification.params}
    });

    if (impl_->transport) {
        impl_->transport->send_message(message);
    }
}

void ServerSession::send_response(const std::string& request_id, const types::ServerResult& result) {
    json::Value message = json::Value::object({
        {"jsonrpc", "2.0"},
        {"id", request_id}
    });

    if (result.is_error()) {
        message["error"] = json::Value::object({
            {"code", static_cast<int>(result.error().code())},
            {"message", result.error().message()}
        });
        if (result.error().data()) {
            message["error"]["data"] = *result.error().data();
        }
    } else {
        message["result"] = result.value();
    }

    if (impl_->transport) {
        impl_->transport->send_message(message);
    }
}

// Server-specific functionality
void ServerSession::send_log_message(types::LogLevel level,
                                   const json::Value& data,
                                   const std::optional<std::string>& logger,
                                   const std::optional<types::RequestId>& related_request_id) {

    auto params = types::LoggingMessageNotificationParams{
        .level = level,
        .data = data,
        .logger = logger
    };

    ServerNotification notification;
    notification.method = "notifications/message";
    notification.params = params.to_json();

    send_notification(notification);
}

void ServerSession::send_resource_updated(const std::string& uri) {
    auto params = types::ResourceUpdatedNotificationParams{
        .uri = uri
    };

    ServerNotification notification;
    notification.method = "notifications/resources/updated";
    notification.params = params.to_json();

    send_notification(notification);
}

std::future<types::CreateMessageResult> ServerSession::create_message(
    const std::vector<types::SamplingMessage>& messages,
    int max_tokens,
    const std::optional<std::string>& system_prompt,
    const std::optional<types::IncludeContext>& include_context,
    const std::optional<double>& temperature,
    const std::optional<std::vector<std::string>>& stop_sequences,
    const std::optional<json::Value>& metadata,
    const std::optional<types::ModelPreferences>& model_preferences,
    const std::optional<types::RequestId>& related_request_id) {

    auto params = types::CreateMessageRequestParams{
        .messages = messages,
        .system_prompt = system_prompt,
        .include_context = include_context,
        .temperature = temperature,
        .max_tokens = max_tokens,
        .stop_sequences = stop_sequences,
        .metadata = metadata,
        .model_preferences = model_preferences
    };

    ServerRequest request;
    request.method = "sampling/createMessage";
    request.params = params.to_json();

    auto future = send_request(request);

    // Convert future<ServerResult> to future<CreateMessageResult>
    auto promise = std::make_shared<std::promise<types::CreateMessageResult>>();
    auto result_future = promise->get_future();

    std::thread([future = std::move(future), promise]() mutable {
        try {
            auto result = future.get();
            if (result.is_error()) {
                promise->set_exception(std::make_exception_ptr(std::runtime_error(result.error().message())));
            } else {
                promise->set_value(types::CreateMessageResult::from_json(result.value()));
            }
        } catch (...) {
            promise->set_exception(std::current_exception());
        }
    }).detach();

    return result_future;
}

std::future<types::ListRootsResult> ServerSession::list_roots() {
    ServerRequest request;
    request.method = "roots/list";
    request.params = json::Value::object({});

    auto future = send_request(request);

    // Convert future<ServerResult> to future<ListRootsResult>
    auto promise = std::make_shared<std::promise<types::ListRootsResult>>();
    auto result_future = promise->get_future();

    std::thread([future = std::move(future), promise]() mutable {
        try {
            auto result = future.get();
            if (result.is_error()) {
                promise->set_exception(std::make_exception_ptr(std::runtime_error(result.error().message())));
            } else {
                promise->set_value(types::ListRootsResult::from_json(result.value()));
            }
        } catch (...) {
            promise->set_exception(std::current_exception());
        }
    }).detach();

    return result_future;
}

std::future<types::ElicitResult> ServerSession::elicit(
    const std::string& message,
    const types::ElicitRequestedSchema& requested_schema,
    const std::optional<types::RequestId>& related_request_id) {

    auto params = types::ElicitRequestParams{
        .message = message,
        .requested_schema = requested_schema
    };

    ServerRequest request;
    request.method = "elicitation/create";
    request.params = params.to_json();

    auto future = send_request(request);

    // Convert future<ServerResult> to future<ElicitResult>
    auto promise = std::make_shared<std::promise<types::ElicitResult>>();
    auto result_future = promise->get_future();

    std::thread([future = std::move(future), promise]() mutable {
        try {
            auto result = future.get();
            if (result.is_error()) {
                promise->set_exception(std::make_exception_ptr(std::runtime_error(result.error().message())));
            } else {
                promise->set_value(types::ElicitResult::from_json(result.value()));
            }
        } catch (...) {
            promise->set_exception(std::current_exception());
        }
    }).detach();

    return result_future;
}

void ServerSession::send_progress_notification(const types::RequestId& progress_token,
                                           double progress,
                                           const std::optional<double>& total,
                                           const std::optional<std::string>& message,
                                           const std::optional<types::RequestId>& related_request_id) {

    auto params = types::ProgressNotificationParams{
        .progress_token = progress_token,
        .progress = progress,
        .total = total,
        .message = message
    };

    ServerNotification notification;
    notification.method = "notifications/progress";
    notification.params = params.to_json();

    send_notification(notification);
}

void ServerSession::send_resource_list_changed() {
    ServerNotification notification;
    notification.method = "notifications/resources/list_changed";
    notification.params = json::Value::object({});

    send_notification(notification);
}

void ServerSession::send_tool_list_changed() {
    ServerNotification notification;
    notification.method = "notifications/tools/list_changed";
    notification.params = json::Value::object({});

    send_notification(notification);
}

void ServerSession::send_prompt_list_changed() {
    ServerNotification notification;
    notification.method = "notifications/prompts/list_changed";
    notification.params = json::Value::object({});

    send_notification(notification);
}

// Resource access
Result<types::ResourceContent> ServerSession::read_resource(const std::string& uri) {
    if (!impl_->resource_manager) {
        return Error::internal_error("No resource manager available");
    }

    // Create a context for this operation
    Context<void> context(shared_from_this(), nullptr, types::RequestId{"resource_read"});

    return impl_->resource_manager->read_resource(uri, context);
}

// Session lifecycle
void ServerSession::initialize(const types::InitializeRequestParams& params) {
    if (impl_->init_state != InitializationState::NotInitialized) {
        throw std::runtime_error("Session already initialized or initializing");
    }

    impl_->init_state = InitializationState::Initializing;
    impl_->client_params = params;

    // Send initialize response
    types::InitializeResult result{
        .protocol_version = "2025-06-18", // Current supported version
        .capabilities = impl_->init_options.capabilities,
        .server_info = types::Implementation{
            .name = impl_->init_options.server_name,
            .version = impl_->init_options.server_version
        },
        .instructions = impl_->init_options.instructions
    };

    send_response(params.request_id, types::ServerResult{result});
}

void ServerSession::initialized() {
    if (impl_->init_state != InitializationState::Initializing) {
        throw std::runtime_error("Session not in initializing state");
    }

    impl_->init_state = InitializationState::Initialized;
}

// Configuration
void ServerSession::set_close_handler(std::function<void()> handler) {
    impl_->close_handler = std::move(handler);
}

// Session information
std::string ServerSession::session_id() const {
    return impl_->session_id;
}

std::chrono::steady_clock::duration ServerSession::uptime() const {
    return std::chrono::steady_clock::now() - impl_->start_time;
}

// Context implementation for server sessions
template<typename ServerSessionT, typename LifespanT>
class ContextImpl {
public:
    ContextImpl(std::weak_ptr<ServerSessionT> session,
               std::shared_ptr<LifespanT> lifespan_context,
               types::RequestId request_id)
        : session_(std::move(session))
        , lifespan_context_(std::move(lifespan_context))
        , request_id_(std::move(request_id)) {}

    // Logging implementation
    void debug(const std::string& message) {
        if (auto session = session_.lock()) {
            session->send_log_message(types::LogLevel::Debug, json::Value(message));
        }
    }

    void info(const std::string& message) {
        if (auto session = session_.lock()) {
            session->send_log_message(types::LogLevel::Info, json::Value(message));
        }
    }

    void warning(const std::string& message) {
        if (auto session = session_.lock()) {
            session->send_log_message(types::LogLevel::Warning, json::Value(message));
        }
    }

    void error(const std::string& message) {
        if (auto session = session_.lock()) {
            session->send_log_message(types::LogLevel::Error, json::Value(message));
        }
    }

    // Progress reporting implementation
    void report_progress(double progress,
                        std::optional<double> total,
                        std::optional<std::string> message) {
        if (auto session = session_.lock()) {
            session->send_progress_notification(request_id_, progress, total, std::move(message));
        }
    }

    // Resource access implementation
    Result<types::ResourceContent> read_resource(const std::string& uri) {
        if (auto session = session_.lock()) {
            return session->read_resource(uri);
        }
        return Error::internal_error("Session expired");
    }

    // User interaction implementation
    template<typename SchemaT>
    Result<types::ElicitationResult> elicit(const std::string& message, const SchemaT& schema) {
        if (auto session = session_.lock()) {
            // This is a simplified implementation
            return Error::unsupported_operation("Elicitation not implemented");
        }
        return Error::internal_error("Session expired");
    }

    // LLM sampling implementation
    Result<types::SamplingResult> create_message(
        const std::vector<types::SamplingMessage>& messages,
        const types::SamplingOptions& options) {
        if (auto session = session_.lock()) {
            // This is a simplified implementation
            return Error::unsupported_operation("LLM sampling not implemented");
        }
        return Error::internal_error("Session expired");
    }

    // Context access implementation
    ServerSessionT& session() {
        if (auto session = session_.lock()) {
            return *session;
        }
        throw std::runtime_error("Session expired");
    }

    LifespanT& lifespan_context() {
        if (!lifespan_context_) {
            throw std::runtime_error("No lifespan context available");
        }
        return *lifespan_context_;
    }

private:
    std::weak_ptr<ServerSessionT> session_;
    std::shared_ptr<LifespanT> lifespan_context_;
    types::RequestId request_id_;
};

} // namespace server
} // namespace mcp